package threadPool;

import cn.hutool.core.date.DateUtil;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.beans.factory.DisposableBean;
import trace.ThreadMdcUtil;

import java.util.LinkedHashMap;
import java.util.concurrent.*;

/**
 * @author lyh
 */
@Slf4j
public class ThreadPoolExecutorMdcWrapper extends ThreadPoolExecutor implements DisposableBean {
    public String poolName;
    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, MyThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
        poolName = threadFactory.getName();
        ThreadMonitor.registerThreadNamePool(poolName);

    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, MyThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        poolName = threadFactory.getName();
        ThreadMonitor.registerThreadNamePool(poolName);

    }

    @Override
    public void execute(Runnable task) {
        if(log.isDebugEnabled()) {
            monitor();
        }
        super.execute(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
        if(log.isDebugEnabled()) {
            monitor();
        }
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        if(log.isDebugEnabled()) {
            monitor();
        }
        return super.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()), result);
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        if(log.isDebugEnabled()) {
            monitor();
        }
        return super.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public Future<?> submit(Runnable task) {
        if(log.isDebugEnabled()) {
            monitor();
        }
        return super.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }


    public LinkedHashMap<String,Object> monitor(){
        LinkedHashMap<String,Object> monitor = new LinkedHashMap<>(8);
        monitor.put("threadPoolName", poolName);
        monitor.put("core", getCorePoolSize());
        monitor.put("getMaximumPoolSize", getMaximumPoolSize());
        monitor.put("poolSize(current)", getPoolSize());
        monitor.put("largestPoolSize(history)", getLargestPoolSize());
        monitor.put("activeCount", getActiveCount());
        monitor.put("queueSize", getQueue().size());
        monitor.put("taskCount", getTaskCount());
        monitor.put("completedTaskCount", getCompletedTaskCount());
        monitor.put("time", DateUtil.now());
        monitor.put("traceId", MDC.get("traceId"));
        log.info("THREAD_POOL_MONITOR:{}", JSONObject.toJSONString(monitor));
        return monitor;

    }
    @Override
    public void destroy() throws Exception {
        super.shutdown();
        log.info("{} shutdown",poolName);
    }
}
